import torch
import numpy as np



def clc_recall(x, label):
    class_num = x.size(1)
    x = torch.argmax(x, dim=1)
    label = torch.argmax(label, dim=1)
    labels_true_num = dict()
    pred_true_num = dict()
    for i in range(class_num):
        tpx = x==i
        tp_pred = label==i
        labels_true_num[str(i)] = torch.sum(tp_pred).item()
        pred_true_num[str(i)] = torch.sum(tpx*tp_pred).item()
    return labels_true_num, pred_true_num

def clc_recall_city(x, label):
    class_num = x.size(1)
    x = torch.argmax(x, dim=1)
    labels_true_num = dict()
    pred_true_num = dict()
    for i in range(class_num):
        tpx = x==i
        tp_pred = label==i
        labels_true_num[str(i)] = torch.sum(tp_pred).item()
        pred_true_num[str(i)] = torch.sum(tpx*tp_pred).item()
    return labels_true_num, pred_true_num

if __name__ == '__main__':
    pred = torch.randint(low=0, high=2, size=(1, 10, 50, 50))
    lb = torch.randint(low=0, high=2, size=(1, 10, 50, 50))
    print(clc_recall(pred, lb))





